import os
import argparse

from concurrent.futures import ThreadPoolExecutor


def run_single_inference(device, model_path, input_path, seed_path, input_name, seed_name, output_path):
    os.system("python3 execute_cann_model.py --device=%d --model_path=%s --input_path=%s --seed_path=%s"
              " --input_name=%s --seed_name=%s --output_path=%s" %
              (device, model_path, input_path, seed_path, input_name, seed_name, output_path))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', type=int, default=0)
    parser.add_argument('--model_path', type=str,
                        default="/root/AscendAtDL2PHY1/models/om")
    parser.add_argument('--input_path', type=str,
                        default="/root/AscendAtDL2PHY1/dataset/tf_testing/input")
    parser.add_argument('--seed_path', type=str,
                        default="/root/AscendAtDL2PHY1/dataset/tf_testing/seed")
    parser.add_argument('--output_path', type=str,
                        default="/root/AscendAtDL2PHY1/results/tf2cann/om2")
    args = parser.parse_args()

    all_model_path = args.model_path
    api_list = os.listdir(all_model_path)
    model_path_list = []
    pool = ThreadPoolExecutor(50)
    for api in api_list:
        mid_model_path = os.path.join(all_model_path, api)
        for exp_str in os.listdir(mid_model_path):
            model_path = os.path.join(mid_model_path, exp_str)
            # If There is no om model int this directory
            if len(os.listdir(model_path)) == 0:
                continue
            input_path = args.input_path + "/" + api + "/" + exp_str
            seed_path = args.seed_path + "/" + api + "/" + exp_str
            input_name = api + "_inputs.json"
            seed_name = api + "_seeds.npz"
            output_path = args.output_path + "/" + api + "/" + exp_str
            os.system("mkdir -p %s" % output_path)
            pool.submit(run_single_inference, args.device, model_path, input_path, seed_path,
                        input_name, seed_name, output_path)
